/*
 * Copyright 2025 the original author or authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package tsj.ai.alibaba.data.analysis.node;

import com.alibaba.cloud.ai.graph.GraphResponse;
import com.alibaba.cloud.ai.graph.OverAllState;
import com.alibaba.cloud.ai.graph.action.NodeAction;
import com.alibaba.cloud.ai.graph.streaming.StreamingOutput;
import lombok.AllArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.document.Document;
import org.springframework.stereotype.Component;
import reactor.core.publisher.Flux;
import tsj.ai.alibaba.data.analysis.service.schema.SchemaService;
import tsj.ai.alibaba.data.analysis.util.ChatResponseUtil;
import tsj.ai.alibaba.data.analysis.util.FluxUtil;
import tsj.ai.alibaba.data.analysis.util.StateUtil;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;

import static tsj.ai.alibaba.data.analysis.constant.Constant.*;


/**
 * 模式召回节点，根据关键字和意图检索相关的数据库模式信息。
 *
 * 此节点负责：-根据用户输入检索相关表-根据提取的关键字检索列文档-组织模式信息以供后续处理-在检索过程中提供流式反馈
 */
@Slf4j
@Component
@AllArgsConstructor
public class SchemaRecallNode implements NodeAction {

	private final SchemaService schemaService;

	@Override
	public Map<String, Object> apply(OverAllState state) throws Exception {
		log.info("Entering {} node", this.getClass().getSimpleName());

		String input = StateUtil.getStringValue(state, QUERY_REWRITE_NODE_OUTPUT,
				StateUtil.getStringValue(state, INPUT_KEY));
		List<String> keywords = StateUtil.getListValue(state, KEYWORD_EXTRACT_NODE_OUTPUT);
		String agentId = StateUtil.getStringValue(state, AGENT_ID);

		// 首先执行业务逻辑-立即调用模式信息
		List<Document> tableDocuments;
		List<List<Document>> columnDocumentsByKeywords;

		// 如果agentId存在，请使用特定于代理的搜索，否则使用全局搜索
		if (agentId != null && !agentId.trim().isEmpty()) {
			log.info("为代理使用特定于代理的模式调用: {}", agentId);
			tableDocuments = new ArrayList<>(schemaService.getTableDocumentsForAgent(agentId, input));
			columnDocumentsByKeywords = schemaService.getColumnDocumentsByKeywordsForAgent(agentId, keywords);
		}
		else {
			log.info("使用全局模式调用（未提供agentId）");
			tableDocuments = schemaService.getTableDocuments(input);
			columnDocumentsByKeywords = schemaService.getColumnDocumentsByKeywords(keywords);
		}

		log.info("表文档数: {}, 关键字关联的列文档: {}", tableDocuments.size(), columnDocumentsByKeywords.size());
		for(Document document: tableDocuments){
			log.info("表文档数: {}", document.toString());
		}
		for(List<Document> document1: columnDocumentsByKeywords){
			for(Document document2: document1){
				log.info("关键字关联的列文档: {}", document2.toString());
			}
		}

		Flux<ChatResponse> displayFlux = Flux.create(emitter -> {
			emitter.next(ChatResponseUtil.createResponse("开始召回Schema信息..."));
			emitter.next(ChatResponseUtil.createResponse("表信息召回完成，数量: " + tableDocuments.size()));
			emitter.next(ChatResponseUtil.createResponse("列信息召回完成，数量: " + columnDocumentsByKeywords.size()));
			emitter.next(ChatResponseUtil.createResponse("Schema信息召回完成."));
			emitter.complete();
		});

		Flux<GraphResponse<StreamingOutput>> generator = FluxUtil.createStreamingGeneratorWithMessages(this.getClass(),
				state, currentState -> {
					log.info("表文档详细信息: {}", tableDocuments);
					log.info("关键字相关列文档详细信息: {}", columnDocumentsByKeywords);
					return Map.of(TABLE_DOCUMENTS_FOR_SCHEMA_OUTPUT, tableDocuments,
							COLUMN_DOCUMENTS_BY_KEYWORDS_OUTPUT, columnDocumentsByKeywords);
				}, displayFlux);


		return Map.of(SCHEMA_RECALL_NODE_OUTPUT, generator);
	}

}
